feat: SimulatorImaging(use_jax=True) + xp-aware preprocess noise - #335
Merged
Conversation
Adds use_jax constructor flag to aa.SimulatorImaging and threads xp through via_image_from. When use_jax=True, via_image_from defaults xp=jnp via a new _xp property; the simulator's Poisson noise generation routes through jax.random.PRNGKey(noise_seed) instead of numpy's RNG. Also adds xp=np parameter to preprocess.poisson_noise_via_data_eps_from and preprocess.data_eps_with_poisson_noise_added so the JAX path routes through jax.random.poisson. The Python `if xp.isnan(noise_map.array).any():` runtime guard is now NumPy- side only — under JAX it would trigger TracerBoolConversionError. JAX users must confirm their exposure_time / background_sky_level are large enough themselves; JAX silently propagates NaN through the trace. Eager JAX path works: simulator.via_image_from with use_jax=True returns an Imaging with jax.Array data. @jax.jit wrapping is currently blocked by an unrelated PyAutoArray limitation: Array2D.native is not jit-traceable because array_2d_via_indexes_from uses indexed assignment of native-index tuples. A separate PyAutoArray task is needed to refactor the slim/native reshape path. The simulator docstring flags this. Part of Phase 2 PR 2 of z_features/jax_user_intro.md. Companion changes: PyAutoLens + PyAutoGalaxy simulator subclass overrides thread xp through via_tracer_from / via_galaxies_from. Design doc: admin_jammy/notes/jax_interface.md Issue: PyAutoArray#334 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
PR 2 of Phase 2 (
z_features/jax_user_intro.md). Addsuse_jax=Truetoaa.SimulatorImagingand routes the Poisson noise pipeline throughjax.randomon the JAX path.Eager JAX path works (
simulator.via_image_fromwithuse_jax=TruereturnsImagingwithjax.Arraydata — useful for batched simulations and parameter sweeps).@jax.jitwrapping is currently blocked by an unrelatedArray2D.nativejit-incompatibility in autoarray's slim/native reshape machinery — flagged in the simulator docstring; needs a separate PR.Companion changes in PyAutoLens (
autolens/imaging/simulator.py) and PyAutoGalaxy (autogalaxy/imaging/simulator.py) — those PRs forwardxpfrom the new_xpproperty on the parent.API Changes
aa.SimulatorImaging(..., use_jax=False)constructor flag.aa.SimulatorImaging._xpproperty (returnsjnpifuse_jax, elsenp).aa.SimulatorImaging.via_image_from(xp=None)— defaults toself._xpwhenNone.preprocess.poisson_noise_via_data_eps_from(..., xp=np)andpreprocess.data_eps_with_poisson_noise_added(..., xp=np)gainxpparameter; JAX path usesjax.random.PRNGKey(noise_seed)+jax.random.poisson.if xp.isnan(noise_map.array).any():NaN-guard now only fires on the NumPy path (Pythonifon a JAX tracer triggersTracerBoolConversionError).See full details below.
Test Plan
test_autoarray/dataset/imaging/test_simulator_use_jax.py— constructor wiring,_xpproperty,via_image_from(xp=None)auto-fallback. Numpy-only assertions per[[feedback_no_jax_in_unit_tests]].test_simulator.py).autolens_workspace_test/scripts/imaging/simulator_use_jax_parity.pyconfirms eager NumPy and JAX paths produce identical simulated images toatol=1e-8for a noise-free Sersic + Isothermal lens. (Ships with the workspace_test follow-up.)Full API Changes (for automation & release notes)
Added
aa.SimulatorImaging(use_jax: bool = False)constructor parameter.aa.SimulatorImaging._xpproperty — returnsjax.numpyifself.use_jax, elsenumpy.Changed signature
aa.SimulatorImaging.via_image_from(image, over_sample_size=None, xp=None)—xpnow defaults toNone, meaning "fall back toself._xp". Explicitxp=nporxp=jnpstill honoured.preprocess.poisson_noise_via_data_eps_from(data_eps, exposure_time_map, seed=-1, xp=np)— newxpparameter.preprocess.data_eps_with_poisson_noise_added(data_eps, exposure_time_map, seed=-1, xp=np)— newxpparameter.Changed behaviour
aa.SimulatorImaging.via_image_from: the NaN-noise-map runtime guard now only runs on the NumPy path. Under JAX it would triggerTracerBoolConversionErrorbecause Pythonif <tracer>:is not allowed. JAX users must confirmexposure_time/background_sky_levelare large enough themselves.preprocess.poisson_noise_via_data_eps_fromon the JAX path: usesjax.random.PRNGKey(seed)+jax.random.poissoninstead ofnp.random.poisson.seed=-1(random per call) derives a time-based PRNGKey.Migration
xp=npor noxp=continue to work unchanged —use_jaxdefaults toFalse.Known limitations
@jax.jitwrap ofvia_image_fromis currently blocked by a pre-existing autoarray limitation:Array2D.nativeroutes througharray_2d_via_indexes_from(indexed assignment of native-index tuples) which is not jit-traceable. Needs a separate slim/native reshape refactor PR.Out of scope (separate PRs)
aa.SimulatorInterferometer.use_jax=True+ signature symmetry fix (Phase 2 PR 3).xp=np+ jnp-backed-grid mismatchValueErrorinAbstractMaker.__init__(Phase 2 PR 4).Array2D.nativejit-safety refactor (separate task).🤖 Generated with Claude Code